import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime
from sqlalchemy import create_engine, String, Date, DECIMAL, BigInteger
import requests
import time
from functools import wraps
from db_config import DBConfig

# 简单的重试装饰器实现
def retry(max_attempts=5, delay=1, backoff=2, exceptions=(Exception,)):
    """
    简单的重试装饰器
    :param max_attempts: 最大重试次数
    :param delay: 初始延迟时间（秒）
    :param backoff: 退避系数
    :param exceptions: 需要重试的异常类型
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            attempts = 0
            current_delay = delay
            while attempts < max_attempts:
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    attempts += 1
                    if attempts == max_attempts:
                        print(f"函数 {func.__name__} 重试 {max_attempts} 次后仍然失败: {e}")
                        raise
                    print(f"函数 {func.__name__} 第 {attempts} 次失败: {e}, {current_delay}秒后重试...")
                    time.sleep(current_delay)
                    current_delay *= backoff  # 指数退避
        return wrapper
    return decorator

def sync_stock_day_to_db(stock_code, stock_name, start_date=None, end_date=None):
    """
    同步股票日线数据到数据库
    """
    end_date = end_date or datetime.now().strftime("%Y%m%d")
    start_date = start_date or end_date
    try:
        raw_df = get_stock_day_data(stock_code, start_date, end_date)
        if raw_df.empty:
            return {
                "success": True,
                "message": f"{stock_name}({stock_code}) 在 {start_date}~{end_date} 未取到数据",
                "records": 0
            }
        df = prepare_stock_day_data(raw_df, stock_code, stock_name)
        if df.empty:
            return {
                "success": True,
                "message": f"{stock_name}({stock_code}) 有数据但缺少昨日收盘价或计算异常，未插入",
                "records": 0
            }
        records = insert_stock_day_to_mysql(df)
        return {
            "success": True,
            "message": f"成功同步 {stock_name}({stock_code}) 数据 {records} 条",
            "records": records
        }
    except Exception as e:
        return {
            "success": False,
            "message": f"数据处理失败: {e}",
            "records": 0
        }

@retry(max_attempts=5, delay=1, backoff=2, exceptions=(requests.RequestException, TimeoutError, ConnectionError))
def get_stock_day_data(symbol, start_date, end_date):
    """
    获取日线数据（带重试机制）：
    先尝试 ak.stock_zh_a_daily，
    如果没有数据或接口失败则用 ak.stock_zh_a_hist 补充。
    """
    df = pd.DataFrame()
    try:
        df = ak.stock_zh_a_daily(
            symbol=symbol,
            start_date=start_date,
            end_date=end_date,
            adjust="qfq"
        )
        if df is not None and not df.empty:
            print(f"成功通过 ak.stock_zh_a_daily 获取 {symbol} 数据，共 {len(df)} 条")
            return df
    except Exception as e:
        print(f"ak.stock_zh_a_daily 调用异常: {e}，尝试使用 ak.stock_zh_a_hist")

    try:
        code_num = symbol.replace("sz", "").replace("sh", "")
        hist_df = ak.stock_zh_a_hist(
            symbol=code_num,
            period="daily",
            start_date=start_date,
            end_date=end_date,
            adjust="qfq"
        )
        if hist_df is not None and not hist_df.empty:
            print(f"成功通过 ak.stock_zh_a_hist 获取 {symbol} 数据，共 {len(hist_df)} 条")
            return convert_hist_to_daily_format(hist_df)
        else:
            print(f"ak.stock_zh_a_hist 未获取到 {symbol} 数据")
    except Exception as e:
        print(f"ak.stock_zh_a_hist 调用异常: {e}")

    return pd.DataFrame()

def convert_hist_to_daily_format(hist_df):
    """
    将 ak.stock_zh_a_hist 的格式转换为与 ak.stock_zh_a_daily 一致
    """
    df = hist_df.rename(columns={
        '日期': 'date',
        '开盘': 'open',
        '收盘': 'close',
        '最高': 'high',
        '最低': 'low',
        '成交量': 'volume',
        '成交额': 'amount',
        '换手率': 'turnover'
    }).copy()
    df['outstanding_share'] = None
    return df[['date', 'open', 'high', 'low', 'close', 'volume', 'amount', 'outstanding_share', 'turnover']]

def prepare_stock_day_data(raw_df, stock_code, stock_name):
    """
    按 stock_day 表字段要求预处理数据，避免 inf 或 -inf
    """
    expected_columns = ['date', 'open', 'high', 'low', 'close', 'volume', 'amount', 'outstanding_share', 'turnover']
    missing_cols = [c for c in expected_columns if c not in raw_df.columns]
    if missing_cols:
        raise ValueError(f"缺少字段: {missing_cols}")

    df = raw_df.rename(columns={
        'date': 'trade_date',
        'open': 'open_price',
        'close': 'close_price',
        'high': 'highest_price',
        'low': 'lowest_price',
        'volume': 'trade_volume',
        'amount': 'turnover_amount',
        'outstanding_share': 'outstanding_share_volume',
        'turnover': 'turnover_rate'
    }).copy()

    # 转换为数值类型
    price_columns = ['open_price', 'close_price', 'highest_price', 'lowest_price']
    for col in price_columns:
        df[col] = pd.to_numeric(df[col], errors='coerce')
    
    df['trade_volume'] = pd.to_numeric(df['trade_volume'], errors='coerce')
    df['turnover_amount'] = pd.to_numeric(df['turnover_amount'], errors='coerce')
    df['outstanding_share_volume'] = pd.to_numeric(df['outstanding_share_volume'], errors='coerce')
    df['turnover_rate'] = pd.to_numeric(df['turnover_rate'], errors='coerce')

    df['last_close_price'] = df['close_price'].shift(1)
    # 避免 last_close_price = 0 导致计算 inf
    df.loc[df['last_close_price'] == 0, 'last_close_price'] = pd.NA
    df = df.dropna(subset=['last_close_price'])
    if df.empty:
        return pd.DataFrame()

    # 涨跌额、涨跌幅、振幅
    df['price_change_amount'] = df['close_price'] - df['last_close_price']
    df['price_change_percent'] = (df['price_change_amount'] / df['last_close_price']) * 100
    df['amplitude_range'] = (df['highest_price'] - df['lowest_price']) / df['last_close_price'] * 100

    # 去掉无穷大或 NaN
    df = df.replace([float('inf'), float('-inf')], pd.NA)
    df = df.dropna(subset=['price_change_percent', 'amplitude_range'])

    # 换手率百分比（如果源数据已经是百分比则不需要乘以100）
    # 根据实际情况调整，如果源数据是小数形式则乘以100
    if df['turnover_rate'].max() <= 1:  # 假设小于1表示是小数形式
        df['turnover_rate'] = df['turnover_rate'] * 100

    # 四舍五入到指定小数位数（根据新表结构调整）
    df['open_price'] = df['open_price'].round(4)  # 调整为4位小数
    df['close_price'] = df['close_price'].round(4)  # 调整为4位小数
    df['last_close_price'] = df['last_close_price'].round(4)  # 调整为4位小数
    df['highest_price'] = df['highest_price'].round(4)  # 调整为4位小数
    df['lowest_price'] = df['lowest_price'].round(4)  # 调整为4位小数
    df['price_change_amount'] = df['price_change_amount'].round(4)  # 调整为4位小数
    df['price_change_percent'] = df['price_change_percent'].round(2)  # 调整为2位小数
    df['amplitude_range'] = df['amplitude_range'].round(2)  # 调整为2位小数
    df['turnover_rate'] = df['turnover_rate'].round(2)  # 调整为2位小数
    df['turnover_amount'] = df['turnover_amount'].round(2)

    # 整数类型字段
    df['trade_volume'] = df['trade_volume'].round(0).astype(pd.Int64Dtype())
    df['outstanding_share_volume'] = df['outstanding_share_volume'].round(0).astype(pd.Int64Dtype())

    df['stock_code'] = stock_code.replace("sz", "").replace("sh", "")
    df['stock_name'] = stock_name
    df['trade_date'] = pd.to_datetime(df['trade_date']).dt.date  # 改为date类型
    
    # 移除create_time和update_time的手动设置，由数据库自动生成
    df = df[['stock_code', 'stock_name', 'trade_date', 'open_price', 'close_price', 'last_close_price',
               'highest_price', 'lowest_price', 'trade_volume', 'turnover_amount', 'outstanding_share_volume',
               'amplitude_range', 'price_change_amount', 'price_change_percent', 'turnover_rate']]

    return df

@retry(max_attempts=3, delay=1, backoff=2, exceptions=(requests.RequestException, TimeoutError, ConnectionError))
def insert_stock_day_to_mysql(df):
    """
    插入 MySQL - 根据新表结构调整数据类型（带重试机制）
    """
    db_config = DBConfig.config
    engine = create_engine(
        f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
    )
    try:
        result = df.to_sql(
            name='stock_day',
            con=engine,
            if_exists='append',
            index=False,
            chunksize=1000,
            dtype={
                'stock_code': String(10),
                'stock_name': String(10),
                'trade_date': Date(),
                'open_price': DECIMAL(10, 4),  # 调整为4位小数
                'close_price': DECIMAL(10, 4),  # 调整为4位小数
                'last_close_price': DECIMAL(10, 4),  # 调整为4位小数
                'highest_price': DECIMAL(10, 4),  # 调整为4位小数
                'lowest_price': DECIMAL(10, 4),  # 调整为4位小数
                'trade_volume': BigInteger(),
                'turnover_amount': DECIMAL(20, 2),
                'outstanding_share_volume': BigInteger(),
                'amplitude_range': DECIMAL(8, 2),  # 调整为2位小数
                'price_change_amount': DECIMAL(10, 4),  # 调整为4位小数
                'price_change_percent': DECIMAL(8, 2),  # 调整为2位小数
                'turnover_rate': DECIMAL(8, 2)  # 调整为2位小数
            }
        )
        print(f"成功插入 {result} 条数据到数据库")
        return result
    finally:
        engine.dispose()

# 使用示例
if __name__ == "__main__":
    # 单个股票同步示例
    #result = sync_stock_day_to_db("sh600000", "浦发银行", "20240101", "20240131")
    #print(result)
    print('end')
